

def distribute_layers(num_layers, pp_size):
    layer_distribution = {}
    layers_per_node = num_layers // pp_size     # gees: layers per process
    remainder = num_layers % pp_size

    for rank in range(pp_size):
        if rank < remainder:
            layer_distribution[rank] = layers_per_node + 1
        else:
            layer_distribution[rank] = layers_per_node
    return layer_distribution



# With 8 layers, 2 stages, and 2 virtual stages, we want an assignment of
# layers to stages like (each list is a model chunk):
# Stage 0: [0, 1] [4, 5]
# Stage 1: [2, 3] [6, 7]
def distribute_layers_with_vpp(num_layers, pp_size, vpp_size):
    # 每个流水线阶段的模型块列表
    layer_distribution = [[] for _ in range(pp_size)]

    # 每个虚拟阶段所包含的层数
    layers_per_virtual_stage = num_layers // (pp_size * vpp_size)
    remainder = num_layers % (pp_size * vpp_size)

    current_layer = 0

    for vpp_rank in range(vpp_size):
        for pp_rank in range(pp_size):
            # 计算当前虚拟阶段的层数
            num_layers_in_chunk = layers_per_virtual_stage
            if remainder > 0:
                num_layers_in_chunk += 1
                remainder -= 1

            # 生成当前模型块的层列表
            layer_chunk = list(range(current_layer, current_layer + num_layers_in_chunk))
            current_layer += num_layers_in_chunk

            # 将模型块添加到对应的流水线阶段
            layer_distribution[pp_rank].append(layer_chunk)

    return layer_distribution


if __name__ == '__main__':
    layer_distribution = distribute_layers(8, 4)
    print(type(layer_distribution))
    print(layer_distribution)
    print(layer_distribution[0])
    print(layer_distribution[1])



    layer_distribution_with_vpp = distribute_layers_with_vpp(16, 4,2)
    print(type(layer_distribution))
    print(layer_distribution)
    print(layer_distribution[0])
    print(layer_distribution[1])
    pp_rank = 0
    vpp_rank = 1
    print(layer_distribution_with_vpp[pp_rank][vpp_rank])
    current_layers_ids = layer_distribution_with_vpp[pp_rank][vpp_rank]
    for current_layer in current_layers_ids:
        print(current_layer)
